package leetcode

import "math/rand"

// https://leetcode.com/problems/random-pick-with-blacklist/
type Solution struct {
	sz       int
	blackMap map[int]int
}

func Constructor(n int, blacklist []int) Solution {
	blackMap, blen := make(map[int]int), len(blacklist)
	sz := n - blen
	for i, v := range blacklist {
		blackMap[v] = i + 100000
	}
	last := n - 1
	for _, v := range blacklist {
		if v < sz {
			_, isBlack := blackMap[last]
			for isBlack {
				last--
				_, isBlack = blackMap[last]
			}
			blackMap[v] = last
			last--
		}
	}
	return Solution{
		sz:       sz,
		blackMap: blackMap,
	}
}

func (t *Solution) Pick() int {
	r := rand.Intn(t.sz)
	if _, ok := t.blackMap[r]; ok {
		r = t.blackMap[r]
	}
	return r
}

/**
 * Your Solution object will be instantiated and called as such:
 * obj := Constructor(n, blacklist);
 * param_1 := obj.Pick();
 */
